import sys
from tts import main as tts
import json
from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall
import warnings
from operation import router
from tools import *
from LLMservice import LLMService,PrettyMessages
warnings.filterwarnings("ignore", category=DeprecationWarning)
#https://platform.openai.com/docs/overview
from stt import main as stt
class TaskAgent:
    def __init__(self):
        self.messages = []
        self.LLMservice = LLMService()
        self.numCheckMain=0
        self.curTask=""
    def addMsg(self,jsonobj:dict):
        self.messages.append(jsonobj)
        print(PrettyMessages(self.messages))
        sys.stdout.flush()

    def stepPre(self, task: str) -> None:
        # 如下是现有文件目录信息以及根据任务的关键词自动搜索到的相关结果\n{ragstr}.
        self.curTask=task
        self.messages.clear()
        self.numCheckMain=0
        self.addMsg({
            "role":"system",
            "content":f'''你是一个机器人的Agent,你需要调用一系列工具来完成任务。如果你不再调用任何工具，任务将会结束。'''
        })
        # 如果你调用TaskComplete工具，任务将会结束。
# 在结束任务之前，你需要说'正在检查某现象以确认任务是否完成'，然后用户会再次给你屏幕截图和a11yTree信息。然后你要仔细查看确认任务是否完成，如果完成了，最后调用TaskComplete工具和speakParams工具（告诉用户执行完毕）。如果没完成，继续操作手机。
        self.addMsg({
            "role":"user",
            "content":f"Now call functions. Remember, the task is {self.curTask}"
        })
    def checkPre(self):
        msg0=self.messages[0].copy()
        self.messages.clear()
        self.addMsg(msg0)
        self.numCheckMain+=1
        if self.numCheckMain>1:
            return
        self.addMsg({
            "role":"user",
            "content":f"请通过截图仔细检查{self.curTask}是否完成。寻找任务被完整地完成的证据。如果已完成，就不需要再次调用函数了，以示任务完成。如果未完成，请继续操作手机。"
        })
        return self.stepMain()
    def stepMain(self):
        reply = self.LLMservice.funcCallLLM(
            "",
            "",
            AGENT_TOOLS,
            self.messages
        )
        if reply.tool_calls:
            for tool_call in reply.tool_calls:
                name = tool_call.function.name
                args = json.loads(tool_call.function.arguments)
                tool_result = router(name, args,self)
                self.addPntToolRtn(tool_result,tool_call)
            return self.stepMain()
        # return self.checkPre()
        return
    def addPntToolRtn(self,tool_result:dict, tool_call: ChatCompletionMessageToolCall)->None:
        """
        # 注意，只会push toolcall result ，不会push toolcall到self.messages
        执行工具调用并构造响应消息
        
        Args:
            tool_name: 工具名称 
            tool_args: 工具参数
            tool_id: 工具调用ID
        Returns:
            构造的工具响应消息
        """
        tool_result["tool_name"] = tool_call.function.name
        result_to_send = tool_result.copy()
        if "imagestr" in tool_result:
            del result_to_send["imagestr"]
        self.addMsg({
            "role": "tool",
            "tool_call_id": tool_call.id, 
            "content": json.dumps(result_to_send)
        })
        if "imagestr" in tool_result:
            self.addMsg({
                "role":"user",
                "content":[{"type":"image_url","image_url":{"url":f"data:image/png;base64,{tool_result['imagestr']}"}}]
            })
        return 
    def main(self,task:str):
        self.stepPre(task)
        self.stepMain()
if __name__=="__main__":
    agent = TaskAgent()
    while True:
        tts("有什么可以帮助你的吗")
        task=stt()
        agent.main(task)
